import numpy as np
import os
from mindspore import dataset


def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


class Sets:
    def __init__(self, basepath) -> None:
        self.basepth = basepath

        self.basename = "data_batch_"

    def reader(self, idx):
        path = os.path.join(self.basepth, self.basename+str(idx))

        data = unpickle(path)

        return data[b'data'], data[b'labels']

    def test(self):
        data = unpickle(os.path.join(self.basepth, "test_batch"))

        return np.array(data[b'data']), np.array(data[b'labels'])

    def run(self, *idx):
        data = []
        target = []
        for it in idx:
            x, y = self.reader(it)
            data.append(x)
            target.append(y)

        data = np.vstack(data)
        target = np.hstack(target)

        return data, target

    def __call__(self):
        x_train, y_train = self.run(1, 2, 3, 4, 5)
        x_test, y_test = self.test()

        return x_train, x_test, y_train, y_test


class MsSets:
    def __init__(self, *inputs) -> None:
        self.x = inputs[0].astype(np.float32)
        self.y = inputs[1].astype(np.int32)

    def __len__(self):
        return self.x.shape[0]

    def __getitem__(self, idx):
        x = self.x[idx].reshape(3, 32, 32)
        y = self.y[idx]

        return x, y


def generate_dataset(*inputs, **kwargs):

    return dataset.GeneratorDataset(MsSets(*inputs), ['image', 'target'], shuffle=True).batch(kwargs['batch'])


if __name__ == "__main__":
    base = "../cifar-10-batches-py"

    m = Sets(base)

    x_train, x_test, y_train, y_test = m()

    print("train X' shape is {}".format(x_train.shape))
    print("test X' shape is {}".format(x_test.shape))

    print("train y' shape is {}".format(y_train.shape))
    print("test y' shape is {}".format(y_test.shape))
